import vtkmodules.all as vtk
from CloudEngine.core.func.color_luts import ColorType, ColorLookUpTable
from CloudEngine.core.func.color_map import color_map


def create_lut(color_type="", colorLevels=256):
    if color_map.get(color_type, None) is not None:
        ctf = vtk.vtkColorTransferFunction()
        rgb_points = color_map[color_type]['RGBPoints']
        for i in range(0, len(rgb_points), 4):
            ctf.AddRGBPoint(rgb_points[i], rgb_points[i + 1], rgb_points[i + 2],
                            rgb_points[i + 3])

        # 创建一个查找表并应用颜色转换函数
        lut = vtk.vtkLookupTable()
        lut.SetNumberOfTableValues(colorLevels)

        for i in range(colorLevels):
            rgb = list(ctf.GetColor(float(i) / colorLevels)) + [1]
            lut.SetTableValue(i, rgb)
        lut.Build()

        return lut
    elif color_type in list(ColorType.__members__.keys()):
        lut = vtk.vtkLookupTable()
        color_lut = ColorLookUpTable(ColorType[color_type], colorLevels)
        lut.SetNumberOfColors(colorLevels)
        for i, color in enumerate(color_lut.color_use_solution):
            lut.SetTableValue(i, color[0], color[1], color[2], 1.0)
        lut.Build()
        return lut
    else:
        lut = vtk.vtkLookupTable()
        lut.SetNumberOfColors(colorLevels)
        lut.Build()
        return lut
